package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;
import geppetto.cat.models.SubstochasticHMM;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.Random;



public class AlignmentsForMoses {
	String corpusDescription;

	int size, maxSentenceLen;
	String modelDir;
	String modelName;
	String baseDir;
	String decodingType;
	float posterior;
	boolean streamInitialized = false;

	PrintStream[] forwardStream, backwardStream;

	

	BilingualCorpus forwardCorpus, backwardCorpus;



	String symmetrization;
	double epsilon; 
	double slack;
	int maxSteps;
	int maxProjectionSteps;
	boolean projectDecoding;
	public AlignmentsForMoses(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceLen = Integer.parseInt(args[2]);
		modelDir = args[3];
		modelName = args[4];
		baseDir = args[5];
		decodingType = args[6];
		posterior = Float.parseFloat(args[7]);
		symmetrization = args[8];
		
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sentence Len " + maxSentenceLen);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("Output Dir " + baseDir);
		System.out.println("Decoding type " + decodingType);
		System.out.println("Posterior " + posterior);
		System.out.println("---");
		
		if (modelName.equalsIgnoreCase("symmetric") || modelName.equalsIgnoreCase("bijective")){
			epsilon = Double.parseDouble(args[9]);
			slack = Double.parseDouble(args[10]);
			maxSteps = Integer.parseInt(args[11]);
			maxProjectionSteps = Integer.parseInt(args[12]);
			projectDecoding = Boolean.parseBoolean(args[13]);
			System.out.println("epsilon " + epsilon);
			System.out.println("slack " + slack);
			System.out.println("maxSteps " + maxSteps);
			System.out.println("maxProjectIterations " + maxProjectionSteps);
			System.out.println("projectDecoding " + projectDecoding);
		}
		
		forwardStream = new PrintStream[3];
		backwardStream = new PrintStream[3];
	}

	public void initializeCorpus(String corpusDesc, int size)
			throws IOException {
		System.out.println("Loading Corpus");
		forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDesc, size, maxSentenceLen);
		backwardCorpus = forwardCorpus.reverse();
		modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
				+ "/" + size + "/model/";
		System.out.println("Finished Loading Corpus");
	}

	

	void initializeStreams(String modelName, String baseDir,
			String decodingType, String sourcePrefix, String foreignPrefix,
			String corpusName) throws FileNotFoundException, IOException {
		System.out.println("Initializing Streams");
		initializeStream(forwardStream, backwardStream, sourcePrefix,
				foreignPrefix, symmetrization, modelName, baseDir, decodingType,
				corpusName);
	}

	void initializeStream(PrintStream[] forwardStream,
			PrintStream[] backwardStream, String sourcePrefix,
			String foreignPrefix, String alignmentName, String modelName,
			String baseDir, String decodingType, String corpusName)
			throws IOException, FileNotFoundException {
		
		String forwardDirection = baseDir + "/" + modelName + "/" + size + "/"
				+ decodingType + "-" + projectDecoding+ "-" + posterior + "/" 
				+ alignmentName + "/" + corpusName + "/"
				+ sourcePrefix + "-" + foreignPrefix + "/model/";
		String backwardDirection = baseDir + "/" + modelName + "/"  + size + "/"
				+ decodingType + "-" + projectDecoding+ "-" + posterior + "/" 
				+ alignmentName + "/" + corpusName + "/"
				+ foreignPrefix + "-" + sourcePrefix + "/model/";
		StaticTools.createDir(forwardDirection);
		StaticTools.createDir(backwardDirection);
		forwardStream[0] = new PrintStream(new FileOutputStream(
				forwardDirection + "aligned." + alignmentName), true, "UTF-8");
		forwardStream[1] = new PrintStream(new FileOutputStream(
				forwardDirection + "aligned.0." + sourcePrefix), true, "UTF-8");
		forwardStream[2] = new PrintStream(new FileOutputStream(
				forwardDirection + "aligned.0." + foreignPrefix), true, "UTF-8");
		backwardStream[0] = new PrintStream(new FileOutputStream(
				backwardDirection + "aligned." + alignmentName), true, "UTF-8");
		backwardStream[1] = new PrintStream(new FileOutputStream(
				backwardDirection + "aligned.0." + foreignPrefix), true,
				"UTF-8");
		backwardStream[2] = new PrintStream(new FileOutputStream(
				backwardDirection + "aligned.0." + sourcePrefix), true, "UTF-8");
	}

	void saveAlignment(PrintStream[] forwardStream,
			PrintStream[] backwardStream, Alignment forward,
			Alignment backward, BilingualCorpus corpus,
			BilingualCorpus backCorpus, String alignmentName, String modelName,
			String baseDir, String decodingType) throws IOException,
			FileNotFoundException {
		String[] data = forward.toMosesFormat(corpus);
		forwardStream[0].println(data[0]);
		forwardStream[1].println(data[1]);
		forwardStream[2].println(data[2]);
		data = backward.toMosesFormat(backCorpus);
		backwardStream[0].println(data[0]);
		backwardStream[1].println(data[1]);
		backwardStream[2].println(data[2]);
	}

	// TODO Atention the models have to be reversed to be save. This is a bug.
	// Probably in the way we are saving them
	public void symtrizeAndSave(Alignment f, Alignment b) throws IOException {
		Alignment forward = null;
		Alignment backward = null;
		if(symmetrization.equalsIgnoreCase("inter")){
			forward = AlignmentSymetrization.intersect(f, b);
			backward = forward.reverse();
		}else if(symmetrization.equalsIgnoreCase("grow")){
			forward = AlignmentSymetrization.growDiag(f, b);
			backward = AlignmentSymetrization.growDiag(b, f);
		}else if (symmetrization.equalsIgnoreCase("grow-final")){
			forward = AlignmentSymetrization.growDiagFinal(f, b);
			backward = AlignmentSymetrization.growDiagFinal(b, f);
		}else if (symmetrization.equalsIgnoreCase("union")){
			forward = AlignmentSymetrization.union(f, b);
			backward = forward.reverse();
		}else if (symmetrization.equalsIgnoreCase("softInter")){
			forward = AlignmentSymetrization.softIntersection(f, b, posterior);
			backward = forward.reverse();
		}else if (symmetrization.equalsIgnoreCase("softUnion")){
			forward = AlignmentSymetrization.softUnion(f, b, posterior);
			backward = forward.reverse();
		}else{
			System.out.println("Unknown symmetrization.... Leaving");
			System.exit(-1);
		}
		saveAlignment(forwardStream, backwardStream, forward, backward
				, forwardCorpus, backwardCorpus, symmetrization, modelName,
				baseDir, decodingType);		
	}

	public void run() throws IOException {
		initializeCorpus(corpusDescription, size);
		AbstractModel models[] = GenericMethods.initializeModels(forwardCorpus, backwardCorpus, modelName, modelDir, epsilon, slack, maxSteps, maxProjectionSteps);
		AbstractModel forward = models[0];
		AbstractModel backward = models[1];
		initializeStreams(modelName, baseDir, decodingType, forwardCorpus
				.getSourceSufix(), forwardCorpus.getForeignSufix(),
				forwardCorpus.getName());

		if (decodingType.equalsIgnoreCase("viter")) {
			for (int i = 0; i < forwardCorpus
					.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.viterbiAlignment(i,
						BilingualCorpus.TRAIN_CORPUS,projectDecoding,null);
				Alignment b = backward.viterbiAlignment(i,
						BilingualCorpus.TRAIN_CORPUS,projectDecoding,null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("....." + i);
			}
		} else if (decodingType.equalsIgnoreCase("post-aer")) {
			float tresh = forward.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,projectDecoding);
			float tresh2 = backward.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,projectDecoding);
			for (int i = 0; i < forwardCorpus
					.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, tresh,projectDecoding,null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, tresh2,projectDecoding,null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("....." + i);
			}

		} else if (decodingType.equalsIgnoreCase("post")) {
			AlignmentsSet tmpA = forward
					.viterbiAlignments(BilingualCorpus.DEV_CORPUS,projectDecoding);
			System.out.println("viterbi forward  alignment points: "
					+ tmpA.numberOfAlignedPoints());
			double fRecall = AlignmentEvaluator.evaluate(tmpA,
					forward.getCorpus().getAlignments(BilingualCorpus.DEV_CORPUS))
					.getRecall();
			tmpA = backward.viterbiAlignments(BilingualCorpus.DEV_CORPUS,projectDecoding);
			System.out.println("viterbi backward alignment points: "
					+ tmpA.numberOfAlignedPoints());
			double bRecall = AlignmentEvaluator
					.evaluate(
							tmpA,
							backward.getCorpus().getAlignments(
									BilingualCorpus.DEV_CORPUS)).getRecall();
			System.out.println("Want to achieve recall of " + fRecall + " "
					+ bRecall);
			// true => fix precision, vary recall
			float tresh = forward.tuneTreshholdPrecisionRecallFixOtherValue(
					BilingualCorpus.DEV_CORPUS, true, (float) fRecall, 0.01f,projectDecoding);
			tmpA = forward.posteriorAlignments(BilingualCorpus.DEV_CORPUS, tresh,projectDecoding,false);
			System.out.println("posterior forward  alignment points: "
					+ tmpA.numberOfAlignedPoints());
			float treshB = backward.tuneTreshholdPrecisionRecallFixOtherValue(
					BilingualCorpus.DEV_CORPUS, true, (float) bRecall, 0.01f,projectDecoding);
			AlignmentsSet tmpB = backward.posteriorAlignments(
					BilingualCorpus.DEV_CORPUS, treshB,projectDecoding,false);
			System.out.println("posterior backward alignment points: "
					+ tmpB.numberOfAlignedPoints());
			for (int i = 0; i < forwardCorpus
					.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, tresh,projectDecoding,null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, treshB,projectDecoding,null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("......." + i);
			}
		} else if (decodingType.startsWith("post-pts")) {
			// sample ~200 sentences from the training corpus
			System.out.println("Forward" + forward);
			System.out.println("Backward" + backward);
			Random r = new Random(-9119629771948792409L);
			int totNumSent = forwardCorpus.getNumberOfTrainingSentences();
			int[] sample = new int[200];
			for (int i = 0; i < sample.length; i++) {
				sample[i] = r.nextInt(totNumSent);
			}
			// count number of aligned points with viterbi
			int numAlignmentsF = forward.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS,projectDecoding).numberOfAlignedPoints();
			int numAlignmentsB = backward.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS,projectDecoding).numberOfAlignedPoints();
			// tune posterior to get that.
			float tresh = forward.tuneThresholdNumPoints(numAlignmentsF, sample,
					BilingualCorpus.TRAIN_CORPUS,projectDecoding);
			float treshB = backward.tuneThresholdNumPoints(numAlignmentsB, sample,
					BilingualCorpus.TRAIN_CORPUS,projectDecoding);
			for (int i = 0; i < forwardCorpus.getNumberOfTrainingSentences(); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, tresh,projectDecoding, null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, treshB,projectDecoding,null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("....." + i);
			}

		} else if (decodingType.startsWith("post-fix")) {
			for (int i = 0; i < forwardCorpus
					.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, posterior, projectDecoding, null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, posterior, projectDecoding, null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("....." + i);
			}
		} else if (decodingType.startsWith("post-f1")) {
			float forwardThresh = forward.tuneTreshholdF1(
					BilingualCorpus.DEV_CORPUS, posterior, projectDecoding);
			float backwardThresh = backward.tuneTreshholdF1(
					BilingualCorpus.DEV_CORPUS, posterior, projectDecoding);
			for (int i = 0; i < forwardCorpus
					.getNumSentences(BilingualCorpus.TRAIN_CORPUS); i++) {
				Alignment f = forward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, forwardThresh, projectDecoding , null);
				Alignment b = backward.posteriorDecodingAlignment(i,
						BilingualCorpus.TRAIN_CORPUS, backwardThresh, projectDecoding, null);
				symtrizeAndSave(f, b);
				if (i % 1000 == 0)
					System.out.print("....." + i);
			}
		} else {
			System.out.println("Unknowon decoding");
		}
	}

	public void closeStream() {
		System.out.println("Closing Streams");
		for (int i = 0; i < 3; i++) {
			forwardStream[i].close();
			backwardStream[i].close();
		}
	}

	public static void main(String[] args) throws IOException {
		AlignmentsForMoses exp = new AlignmentsForMoses(args);
		exp.run();
		exp.closeStream();
	}

}
